import sqlite3
from datetime import datetime, timedelta
import requests
import pandas as pd
import os
import sys

# 数据库配置
DB_FILE = 'stock_data.db'


def fetch_stock_data(symbol, exchange='SZSE', start_date="2010-01-01", end_date=None,qfqday=True):

    if end_date is None:
        end_date = datetime.now().strftime("%Y-%m-%d")
    
    exchange_code = exchange.lower()
    bar_data = []
    current_date = datetime.strptime(start_date, "%Y-%m-%d")
    end_date = datetime.strptime(end_date, "%Y-%m-%d")
    
    print(f"开始获取 {exchange}{symbol} 数据，时间范围: {start_date} 至 {end_date.date()}")
    
    while current_date <= end_date:
        chunk_end = min(current_date + timedelta(days=365), end_date)
        start_str = current_date.strftime("%Y-%m-%d")
        end_str = chunk_end.strftime("%Y-%m-%d")
        
        print(f"获取数据: {start_str} 至 {end_str}")
        
        if qfqday:
            url = f"http://web.ifzq.gtimg.cn/appstock/app/fqkline/get?param={exchange_code}{symbol},day,{start_str},{end_str},2000,qfqday"
        else:
            url = f"http://web.ifzq.gtimg.cn/appstock/app/fqkline/get?param={exchange_code}{symbol},day,{start_str},{end_str},2000,qfq"
        print(f"url:",url)
        try:
            response = requests.get(url, timeout=10)
            data = response.json()
            
            if not data.get('data') or f"{exchange_code}{symbol}" not in data['data']:
                print(f"获取数据失败，响应: {data}")
                current_date = chunk_end + timedelta(days=1)
                continue
            
            for record in data["data"][f"{exchange_code}{symbol}"]["day"]:
                if len(record) < 6:
                    continue
                if symbol == "null":
                    bar_data.append({
                        'Date': datetime.strptime(record[0], "%Y-%m-%d"),
                        'Open': round(float(record[1])/1000, 2),
                        'High': round(float(record[3])/1000, 2),
                        'Low': round(float(record[4])/1000, 2),
                        'Close': round(float(record[2])/1000, 2),
                        'Adj Close': round(float(record[2])/1000, 2),
                        'Volume': int(float(record[5]))
                    })
                    pass
                else:
                    bar_data.append({
                        'Date': datetime.strptime(record[0], "%Y-%m-%d"),
                        'Open': float(record[1]),
                        'High': float(record[3]),
                        'Low': float(record[4]),
                        'Close': float(record[2]),
                        'Adj Close': float(record[2]),
                        'Volume': int(float(record[5]))
                    })
                
        except Exception as e:
            print(f"获取数据出错: {e}")
            return  bar_data
        
        current_date = chunk_end + timedelta(days=1)
    
    print(f"共获取到 {len(bar_data)} 条数据")
    
    # # 动态添加测试数据（如果获取的数据不包含start_date之后的新数据）
    # start_date_dt = datetime.strptime(start_date, "%Y-%m-%d")
    # # 检查是否有start_date及之后的新数据
    # has_new_data = any(item['Date'] >= start_date_dt for item in bar_data)
    
    # if not has_new_data:
    #     print(f"没有获取到 {start_date} 之后的新数据，将添加测试数据")
    #     # 创建一条新的测试数据，日期为start_date
    #     bar_data.append({
    #         'Date': start_date_dt,
    #         'Open': 4460.32,
    #         'High': 4500.0,
    #         'Low': 4450.0,
    #         'Close': 4480.0,
    #         'Adj Close': 4480.0,
    #         'Volume': 250000000
    #     })
    
    return bar_data


def save_to_csv(data, symbol, exchange):
    """将股票数据保存为CSV文件，支持追加数据"""
    if not data:
        print("没有数据需要保存为CSV文件")
        return
    
    # 创建DataFrame
    df = pd.DataFrame(data)
    
    # 将datetime列格式化为字符串
    df['Date'] = df['Date'].dt.strftime('%Y-%m-%d')
    # 修复Volume列转换错误：先转换为浮点数，再转换为整数
    df['Volume'] = df['Volume'].astype(float).astype(int)
    # 创建保存目录（如果不存在）
    csv_dir = 'csv_data'
    if not os.path.exists(csv_dir):
        os.makedirs(csv_dir)
        print(f"创建目录: {csv_dir}")
    
    # 构建文件名
    file_name = f"{csv_dir}/{exchange}_{symbol}_stock_data.csv"
    
    # 检查文件是否存在
    if os.path.exists(file_name):
        print(f"CSV文件已存在: {file_name}，将尝试追加数据")
        try:
            # 读取已有数据
            existing_df = pd.read_csv(file_name, encoding='utf-8-sig')
            # 合并新数据和已有数据
            combined_df = pd.concat([existing_df, df])
            # 根据Date列去重，保留最后一条记录（如果有重复）
            combined_df.drop_duplicates(subset=['Date'], keep='last', inplace=True)
            # 按日期排序
            combined_df.sort_values('Date', inplace=True)
            # 保存合并后的数据
            combined_df.to_csv(file_name, index=False, encoding='utf-8-sig')
            print(f"数据已追加到CSV文件: {file_name}")
        except Exception as e:
            print(f"追加数据出错: {e}")
            # 出错时仍然尝试保存新数据（覆盖模式）
            df.to_csv(file_name, index=False, encoding='utf-8-sig')
            print(f"数据已保存到CSV文件（覆盖模式）: {file_name}")
    else:
        # 文件不存在，创建新文件
        df.to_csv(file_name, index=False, encoding='utf-8-sig')
        print(f"数据已保存到CSV文件: {file_name}")

def get_last_record_date(symbol, exchange):
    """获取CSV文件中某只股票的最后记录日期"""
    # 构建CSV文件路径
    csv_dir = 'csv_data'
    file_name = f"{csv_dir}/{exchange}_{symbol}_stock_data.csv"
    
    # 检查文件是否存在
    if not os.path.exists(file_name):
        print(f"CSV文件不存在: {file_name}")
        return None
    
    try:
        # 使用pandas读取CSV文件
        df = pd.read_csv(file_name, encoding='utf-8-sig')
        
        # 如果文件为空或只有表头
        if len(df) == 0:
            print(f"CSV文件为空: {file_name}")
            return None
        
        # 获取最后一行的Date列值
        last_record_date_str = df['Date'].iloc[-1]
        print(f"从CSV文件读取的最后记录日期: {last_record_date_str}")
        
        # 转换为datetime对象并返回
        last_record_date = datetime.strptime(last_record_date_str, "%Y-%m-%d")
        return last_record_date
    except Exception as e:
        print(f"读取CSV文件最后记录日期出错: {e}")
        return None


def bk_pull_stock_data(symbol, exchange):
    # 默认参数
    # symbol = '000300'  # 
    # exchange = 'SH'    # 
    # 2. 获取最后记录日期，确定下载范围
    last_date = get_last_record_date(symbol, exchange)
    # 如果有最后记录日期，则从该日期的下一天开始获取数据；否则使用默认日期
    if last_date:
        start_date = (last_date + timedelta(days=1)).strftime("%Y-%m-%d")
        print(f"最后记录日期: {last_date.date()}")
        print(f"将从 {start_date} 开始获取数据")
    else:
        start_date = "2020-01-01"
        print(f"没有找到最后记录日期，将从默认日期 {start_date} 开始获取数据")
    
    # 3. 获取股票数据
    stock_data = fetch_stock_data(symbol, exchange, start_date=start_date, qfqday=True)
    if len(stock_data) == 0:
        print(f"retry : {symbol} {exchange} {start_date}")
        stock_data = fetch_stock_data(symbol, exchange, start_date=start_date, qfqday=False)
        return
    # 打印准备保存的数据信息
    print(f"准备保存的数据: {len(stock_data)} 条，数据样本: {stock_data[:1] if stock_data else '无'}")
    
    # 5. 保存到CSV文件
    save_to_csv(stock_data, symbol, exchange)
    
    

if __name__ == "__main__":
    # 使用sys.argv解析命令行参数
    if len(sys.argv) != 3:
        print("参数，例如\n sh: python pull_real_data_stock.py 000300 SH")
        print("sz: python pull_real_data_stock.py 300498 SZ")
        sys.exit(1)
    symbol = sys.argv[1]
    exchange = sys.argv[2]
    print(f"get stock date: {symbol} {exchange}".format(symbol, exchange))
    bk_pull_stock_data(symbol, exchange)
